#include "log_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include "graph/utils/type_utils.h"
#include<cmath>

namespace optiling {
const uint32_t BLOCK_SIZE = 32;
const uint32_t BUFFER_NUM = 2;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
        LogTilingData tiling;

        uint64_t ubSize;
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
        ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
        auto coreNum = ascendcPlatform.GetCoreNum();
        auto socVersion = ascendcPlatform.GetSocVersion();

        if (socVersion != platform_ascendc::SocVersion::ASCEND910B && socVersion != platform_ascendc::SocVersion::ASCEND310B && context->GetInputDesc(0)->GetDataType() == ge::DT_BF16)
        {
            return ge::GRAPH_FAILED;
        }

        uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
        uint32_t typeLength = 0;
        ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(0)->GetDataType(), typeLength);
        uint32_t inputLength = inputNum * typeLength;
        uint32_t inputBytes = inputLength / inputNum;

        uint32_t ubDataNumber = (context->GetInputDesc(0)->GetDataType() == ge::DT_FLOAT) ? 3 : 5;
        uint32_t tileBlockNum = (ubSize / BLOCK_SIZE / BUFFER_NUM) / ubDataNumber;
        uint32_t tileDataNum = (tileBlockNum * BLOCK_SIZE) / inputBytes;

        uint32_t inputLengthAlgin32 = (((inputLength + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE);
        coreNum = (coreNum < inputLengthAlgin32 / BLOCK_SIZE) ? coreNum : inputLengthAlgin32 / BLOCK_SIZE;
        coreNum = (coreNum >= 1) ? coreNum : 1;
        uint32_t everyCoreInputBlockNum = inputLengthAlgin32 / BLOCK_SIZE / coreNum;
        uint32_t tailBlockNum = (inputLengthAlgin32 / BLOCK_SIZE) % coreNum;

        uint32_t smallCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
        uint32_t smallTileNum = everyCoreInputBlockNum / tileBlockNum;
        uint32_t finalSmallTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? smallTileNum : smallTileNum + 1;
        uint32_t smallTailDataNum = smallCoreDataNum - (tileDataNum * smallTileNum);
        smallTailDataNum = smallTailDataNum == 0 ? tileDataNum : smallTailDataNum;

        everyCoreInputBlockNum += 1;
        uint32_t bigCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
        uint32_t bigTileNum = everyCoreInputBlockNum / tileBlockNum;
        uint32_t finalBigTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? bigTileNum : bigTileNum + 1;
        uint32_t bigTailDataNum = bigCoreDataNum - tileDataNum * bigTileNum;
        bigTailDataNum = bigTailDataNum == 0 ? tileDataNum : bigTailDataNum;

        tiling.set_smallCoreDataNum(smallCoreDataNum);
        tiling.set_bigCoreDataNum(bigCoreDataNum);
        tiling.set_tileDataNum(tileDataNum);
        tiling.set_smallTailDataNum(smallTailDataNum);
        tiling.set_bigTailDataNum(bigTailDataNum);
        tiling.set_finalSmallTileNum(finalSmallTileNum);
        tiling.set_finalBigTileNum(finalBigTileNum);
        tiling.set_tailBlockNum(tailBlockNum);

        // 属性值获取和设置
        const gert::RuntimeAttrs *attrs = context->GetAttrs();
        const float *base = attrs->GetAttrPointer<float>(0);
        if (*base == -1.0f) {
            tiling.set_base(1.0f);
        }
        else if (*base > 0.0f && *base != 1.0f){
            tiling.set_base(1.0f / log(*base));
        }


        const float *scale = attrs->GetAttrPointer<float>(1);
        tiling.set_scale(*scale);
        const float *shift = attrs->GetAttrPointer<float>(2);
        tiling.set_shift(*shift);


        context->SetBlockDim(coreNum);
        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
        size_t *currentWorkspace = context->GetWorkspaceSizes(1);
        currentWorkspace[0] = 0;
        return ge::GRAPH_SUCCESS;
    }
}



namespace ge 
{
    static ge::graphStatus InferShape(gert::InferShapeContext* context)
    {
        // 简单地将输入的形状复制给输出，表示输出与输入形状相同。
        const gert::Shape* x1_shape = context->GetInputShape(0);
        gert::Shape* y_shape = context->GetOutputShape(0);
        *y_shape = *x1_shape;
        return GRAPH_SUCCESS;
    }
}


namespace ops 
{
    class Log : public OpDef 
    {
        public:
            explicit Log(const char* name) : OpDef(name)
            {
                // 输入定义
                this->Input("x")
                    .ParamType(REQUIRED)
                    .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT})
                    .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
                    .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
                // 输出定义
                this->Output("y")
                    .ParamType(REQUIRED)
                    .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT})
                    .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
                    .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

                // 属性定义
                this->Attr("base").AttrType(OPTIONAL).Float(-1.0);
                this->Attr("scale").AttrType(OPTIONAL).Float(1.0);
                this->Attr("shift").AttrType(OPTIONAL).Float(0.0);

                // 形状推导函数
                this->SetInferShape(ge::InferShape);

                // Tiling策略
                this->AICore()
                    .SetTiling(optiling::TilingFunc);
                this->AICore().AddConfig("ascend910b").AddConfig("ascend310p");

            }
    };

    OP_ADD(Log);
}
